# api.py
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from transformers import BertTokenizer, BertForSequenceClassification
import torch
import logging

# 初始化 FastAPI
app = FastAPI()

# 加载模板和静态文件
templates = Jinja2Templates(directory="templates")
app.mount("/static", StaticFiles(directory="static"), name="static")

# 配置日志记录
logger = logging.getLogger()
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(message)s', '%Y-%m-%d %H:%M:%S')
file_handler = logging.FileHandler('logs/detection.log', encoding='utf-8')
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)

# 加载模型和 tokenizer
model_path = "models/procedure_injection_bert"
model = BertForSequenceClassification.from_pretrained(model_path)
tokenizer = BertTokenizer.from_pretrained(model_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# 标签映射
label_map = {
    0: "✅ 正常 SQL",
    1: "🚨 SQL 注入 ⚠️",
    2: "🚨 存储过程注入 ⚠️"
}

# 存储过程关键字
STORED_PROCEDURE_KEYWORDS = [
    "EXEC",
    "EXECUTE",
    "sp_executesql",
    "DECLARE",
    "SET",
    "OPENROWSET",
    "BEGIN",
    "END",
    "CREATE PROCEDURE",
    "ALTER PROCEDURE",
    "DROP PROCEDURE",
    "CREATE FUNCTION",
    "ALTER FUNCTION",
    "DROP FUNCTION"
]

# 清洗文本
def clean_text(text):
    text_cleaned = text
    for keyword in STORED_PROCEDURE_KEYWORDS:
        text_cleaned = text_cleaned.replace(keyword, "")
        text_cleaned = text_cleaned.replace(keyword.lower(), "")
    return text_cleaned.strip()

# API请求模型
class SQLRequest(BaseModel):
    query: str

# 首页
@app.get("/", response_class=HTMLResponse)
def read_root(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

# /predict API
@app.post("/predict")
def predict_sql(req: SQLRequest):
    query = req.query.strip()

    # Step 1: 是否含存储过程关键字
    contains_proc_keyword = any(kw.lower() in query.lower() for kw in STORED_PROCEDURE_KEYWORDS)

    if contains_proc_keyword:
        # Step 2: 去除关键字
        cleaned_query = clean_text(query)

        # Step 3: 剩余为空 / 纯数字 → 正常
        if cleaned_query == "" or cleaned_query.replace(" ", "").isdigit():
            final_label = 0
        else:
            # 仍有内容，走模型判断
            inputs = tokenizer(cleaned_query, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
            with torch.no_grad():
                outputs = model(**inputs)
                pred_cleaned = torch.argmax(outputs.logits, dim=1).item()

            if pred_cleaned == 1:  # 剩余仍是注入
                final_label = 2
            else:
                final_label = 0  # 合法调用

    else:
        # 普通流程直接判断
        inputs = tokenizer(query, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
        with torch.no_grad():
            outputs = model(**inputs)
            final_label = torch.argmax(outputs.logits, dim=1).item()

    # 输出结果
    result = label_map.get(final_label, "未知类别")

    # 记录日志
    logging.info(f"输入 SQL: {query}, 检测结果: {result}")

    return {"label": result}

# /logs API：返回日志文件内容
@app.get("/logs", response_class=PlainTextResponse)
def get_logs():
    try:
        with open("logs/detection.log", "r", encoding="utf-8") as f:
            log_content = f.read()
        return log_content
    except Exception as e:
        return f"无法读取日志文件: {str(e)}"
